import numpy as np
import re
import matplotlib.pyplot as plt
from matplotlib import rc
from os import path
from utils import count_lines
from collections import defaultdict

"""
Printable label and description of the different parameters and measures.
"""
descriptions = {
        'n_positive': ('Training pos size', 'Number of positive examples used in training'),
        'n_negative': ('Training neg size', 'Number of negative examples used in training'),
        'n_test': ('Test size', 'Test input size in number of sentences'),
        'threshold': ('Threshold', 'Decision probability threshold of the classifier'),
        'test_corpus': ('Test corpus', 'Corpus used for testing'),
        'train_corpus': ('Training corpus', 'Corpus used for training'),
        'filtering': ('Candidate filtering', 'Candidate filtering'),
        'min_overlap': ('Filter threshold', 'Threshold parameter of the word-overlap filter'),
        'length_ratio': ('Filter length ratio', 'Maximum length ratio used in the filter'),
        'noise_level': ('Noise level', 'Noise level in percentage'),
        'giza_model_path': ('GIZA model', 'GIZA model'),
        'recall_filter': ('Filter recall', ''),
        'precision_filter': ('Filter precision', ''),
        'recall_prior': ('Recall prior', ''),
        'precision_prior': ('Precision prior', ''),
        'recall': ('Recall', ''),
        'precision': ('Precision', ''),
        'f1': ('F-measure', ''),
        'f1_prior': ('F1 prior', ''),
        'filter_time': ('Filter time', ''),
        'time': ('Total time', ''),
        'classifier_time': ('Classifier time', ''),
        'filter_drop': ('Filter drop rate', ''),
        'f1_yasa': ('YASA F1', ''),
        'recall_yasa': ('YASA recall', ''),
        'precision_yasa': ('YASA precision', '')
        }

plot_markers = ['o', 's', '^', 'x', '+']
plot_colors = ['0.7', '0.3', 'k', '0.4', 'm', 'c', 'y', 'g']

def read_experiment_result(filename, x_label, y_labels=None):
    res = dict()

    x_label = x_label.lower()
    if y_labels:
        y_labels = [y_label.lower() for y_label in y_labels]
    
    with open(filename) as input_file:
        parameters = None

        for line in input_file:
            if line.startswith('#'): 
                parameters = dict((k.lower(), v) for k, v in re.findall(r'(\w+)=([^,\s]*)', line))
            elif parameters is not None and x_label in parameters:
                x = parameters[x_label]

                # Match only booleans, integers and floats
                if re.match('(True|False|\d+|\d+\.\d*|\d*\.\d+)$', x):
                    x = eval(x)

                statistics = dict()
                statistics = dict((k.lower(), float(v)) for k, v in re.findall(r'(\w+)=([^,]*)', line)
                        if y_labels is None or k.lower() in y_labels)

                res[x] = statistics
    return res

def plot_file(input_filename, output_filename, x_label, y_labels=None, y_label=None, title=None, bar=False, y_margin=0, x_values_mapping=[], alpha=1, leg_loc='upper left'):
    res = read_experiment_result(input_filename, x_label, y_labels)
    res = dict((path.basename(x) if isinstance(x, str) else x, y) for (x, y) in res.iteritems())

    Y = defaultdict(list)
    X = []
    for x_value, y_values in sort_items(res, x_values_mapping):
        X.append(x_value)
        for y_label_, y_value in y_values.items():
            if y_labels:
                y_label_nice = y_labels[y_label_]
            else:
                y_label_nice = descriptions[y_label_][0]
            Y[y_label_nice].append(y_value)
    
    x_label_nice = descriptions[x_label][1]
    plot(X, Y, x_label_nice, output_filename, y_label, title, bar, y_margin, alpha=alpha, leg_loc=leg_loc)

def sort_items(dictionary, x_values_mapping=[]):
    """ ``x_values_mapping'' specifies both a mapping for the keys of the dictionary, and a sort order. """
    items = dictionary.items()
    if not x_values_mapping:
        return sorted(items)
    else:
        x_order, _ = zip(*x_values_mapping)
        mapping = dict(x_values_mapping)
        items = sorted(items, key=lambda p: x_order.index(p[0]))
        items = [(mapping[x] if x in mapping else x, y) for x, y in items]
        return items

def plot_files(filenames, suffixes, output_filename, x_label, y_labels=None, y_label=None, title=None, bar=False,
        x_values_mapping=[]):
    """
    Merge several 2D experiments in a single plot.
    If specified, ``y_labels'' must be a dictionary, whose keys are the labels of the y values to keep. The values of the dictionary specify the text to write in the legend (+ file suffix). If ``y_labels'' is None, all values are kept, and the default descriptions are used in the legend.
    ``suffixes'' is a list containing the suffix to append to the legend, for each file.
    """
    results = [read_experiment_result(filename, x_label, y_labels) for filename in filenames]
    results_ = defaultdict(list)

    for result in results:
        for x, statistics in result.items():
            if isinstance(x, str):
                x = path.basename(x)
            results_[x].append(statistics)
    results = dict((x, l) for x, l in results_.items() if len(l) == len(results))

    X, all_Y = zip(*sort_items(results, x_values_mapping))

    values = defaultdict(list)
    for Y in all_Y:
        for suffix, y_values in zip(suffixes, Y):
            for y_label_, y_value in y_values.items():
                if y_labels:
                    y_label_nice = y_labels[y_label_]
                else:
                    y_label_nice = descriptions[y_label_][0]
                values[y_label_nice + suffix].append(y_value)

    x_label_nice = descriptions[x_label][1]
    plot(X, values, x_label_nice, output_filename, y_label, title, bar)

def plot_values(X, Y, x_label, output_filename, title=None):
    """ Same as ``plot'', but replaces the labels by their description. """
    Y = dict((descriptions[y_label][0], y_values) for y_label, y_values in Y.items())
    x_label_nice = descriptions[x_label][1] if x_label in descriptions else x_label.capitalize()
    plot(X, Y, x_label_nice, output_filename, title)    

def plot(X, Y, x_label, output_filename, y_label=None, title=None, bar=False, y_margin=0, alpha=0.5, leg_loc='upper left', linestyle='solid', reverse_x=False):
    """
    X is a list of x values
    Y is a list of pairs (y_label, y_values)
    """
    rc('text', usetex=True)

    fig = plt.figure(figsize=(6, 4))
    ax = fig.add_subplot(111)
    y_margin += 0.1
    ax.set_position([0.1, y_margin, 0.85, 0.95 - y_margin])

    bar = bar or any(x for x in X if isinstance(x, str) or isinstance(x, bool))
    X = [x.replace('_', '\\_') if isinstance(x, str) else x for x in X] # For LaTeX

    i = 0
    w = 1.0 / (len(Y) + 1)
    ind = np.arange(len(X))

    for c, m, (k, array) in zip(plot_colors, plot_markers, sorted(Y.items())):
        label = k

        if bar:
            ax.bar(ind + w * i, array, w, color=c, label=label)
        else:
            ax.plot(X, array, marker=m, color=c, linestyle=linestyle, label=label)
        
        i += 1
    
    if bar:
        ax.set_xticks(ind + w*(len(Y) / 2))
        ax.set_xticklabels(X)
        low = min(min(array) for array in Y.values())
        high = max(max(array) for array in Y.values())
        plt.ylim(ymin=max(0, low - 0.2 * (high - low)))

    if title is not None:
        plt.title(title) 
    plt.xlabel(x_label)
    if reverse_x:
        plt.gca().invert_xaxis()
    if y_label is not None:
        plt.ylabel(y_label)

    leg = ax.legend(loc=leg_loc, prop={'size': 12})
    leg.get_frame().set_alpha(alpha)
    fig.savefig(output_filename + '.pdf')

